Задано
подвешенное дерево, содержащее n (1
≤ n ≤ 106) вершин. Каждая
вершина покрашена в один из n цветов.
Требуется для каждой вершины v
вычислить количество различных цветов, встречающихся в поддереве с корнем v.
Вход. В первой строке задано число n. Следующие n строк
описывают вершины по одной в строке. Описание очередной вершины i имеет вид pi ci, где pi – номер родителя вершины i, а ci – цвет
вершины i (1 ≤ ci ≤ n). Для корня дерева pi = 0.
Выход. Выведите n чисел, обозначающих количества различных цветов в поддеревьях с
корнями в вершинах 1, 2, ..., n.
Пример
входа |
Пример
выхода |
5 2 1 3 2 0 3 3 3 2 1 |
1 2 3 1 1 |
РЕШЕНИЕ
поиск в глубину
Анализ алгоритма
Запустим поиск в
глубину из корня дерева. Для каждой вершины i
храним множество si, в
котором будем накапливать цвета вершин в ее поддеревьях. Если j – сын вершины i при поиске в глубину, то sj
должно включаться в si.
Количество различных цветов в поддереве с корнем i равно размеру множества si.
Пример
Слева возле каждой вершины приведен
ее цвет. Справа возле каждой вершины приведено множество цветов в поддереве с
корнем в этой вершине.
Реализация алгоритма
В ячейке color[i] храним цвет i-ой вершины. Во множестве s[i]
будем накапливать цвета в поддереве с корнем i. Ориентированное дерево храним в списке смежности графа g. В res[i] храним количество различных цветов в
поддереве с корнем i.
#define MAX 1000010
int color[MAX], res[MAX];
set<int>
s[MAX];
vector<vector<int> > g;
Функция dfs совершает поиск в глубину из вершины v. Изначально заносим в s[v]
цвет вершины v. Для каждого ребра
дерева (v, to) добавляем множество s[to]
в s[v]. Количество различных цветов в
поддереве с корнем v равно размеру
множества s[v], заносим его в res[v].
void dfs(int v)
{
int i, to;
s[v].insert(color[v]);
for(i = 0; i < g[v].size(); i++)
{
to = g[v][i];
dfs(to);
Если размер множества s[v] меньше размера
множества s[to], то меняем их местами. Далее содержимое
меньшего множества s[to] добавляем во
множество s[v].
if (s[v].size() <
s[to].size()) s[v].swap(s[to]);
s[v].insert(s[to].begin(), s[to].end());
Очищаем множество s[to] – оно нам больше
не пригодится.
s[to].clear();
}
res[v] = s[v].size();
}
Основная часть программы. Читаем входные данные.
scanf("%d",&n);
g.resize(n+1);
for(i = 1; i <= n; i++)
{
scanf("%d %d",&p,&c);
g[p].push_back(i);
color[i] = c;
}
Запускаем поиск в глубину из корня дерева – нулевой вершины.
dfs(0);
Выводим ответ.
for(i = 1; i <= n; i++)
printf("%d ",res[i]);
printf("\n");